#!/usr/bin/env python
# -*- coding: utf-8 -*-

import envs
import genpy
import rospy
import actionlib

from std_msgs.msg import Float64
from control_msgs.msg import JointTrajectoryControllerState
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint

from npu_robot.msg import ArmSwingAction, ArmSwingFeedback, ArmSwingResult

# 设置相机分辨率
CAMERA_WIDTH = 640
CAMERA_HEIGHT = 480

# 偏差阈值
BIAS_THRESHOLD = 0.1
# 控制信号比例
CONTROL_RATIO = 0.5


class ArmSwingActionServer():
    # 创建用于发布反馈(feedback)/结果(result)的消息
    _feedback = ArmSwingFeedback()
    _result = ArmSwingResult()

    def __init__(self):
        # 机械臂摇摆服务端
        self._action_server = actionlib.SimpleActionServer('xarm_controller/arm_swing', ArmSwingAction, execute_cb=self.execute_callback, auto_start=False)

        # 物体中心点横坐标订阅者
        self.object_center_x_sub = rospy.Subscriber('/camera/image_raw/object_center_x', Float64, self.object_center_x_callback, queue_size=1)

        # 机械臂状态发布者
        self.joint_cmd_pub = rospy.Publisher('/xarm_controller/command', JointTrajectory, queue_size=10)

        # 机械臂状态订阅者
        self.joint_state_sub = rospy.Subscriber('/xarm_controller/state', JointTrajectoryControllerState, self.joint_state_callback, queue_size=10)

        # 机械臂关节初始化
        self.joint_states = []

        # 机械臂关节控制初始化
        self.joint_cmd_data = JointTrajectory()
        self.joint_cmd_data.joint_names = ['arm_joint_1', 'arm_joint_2', 'arm_joint_3', 'arm_joint_4', 'arm_joint_5']
        self.joint_cmd_data_point = JointTrajectoryPoint()
        self.joint_cmd_data_point.time_from_start = genpy.Duration(0.5)

        # 物体中心点回调次数
        self.object_center_x_callback_times = 0
        # 物体中心点回调次数阈值
        self.object_center_x_callback_times_threshold = 3

        # 机械臂到达偏差阈值之内的次数
        self.bias_x_correct_times = 0
        # 机械臂到达偏差阈值之内的次数阈值
        self.bias_x_correct_times_threshold = 50

        # 标志位 - action是否完成
        self.flag = False

        # 启动服务端
        self._action_server.start()

        rospy.spin()

    def execute_callback(self, goal):
        """
        action请求回调
        """
        rospy.logdebug('arm swing action execute callback...')

        while not self.flag:
            pass

        self.flag = False

    def joint_state_callback(self, req):
        """
        机械臂状态回调函数
        """
        rospy.logdebug('joint state callback...')
        rospy.logdebug('req.actual.positions: %s', req.actual.positions)

        self.joint_states = list(req.actual.positions)

        rospy.logdebug('joint state callback OK')

    def object_center_x_callback(self, data):
        """
        物体中心点横坐标回调函数
        """
        rospy.logdebug('object center x callback...')

        # 确保行为还没有被取消
        if not self._action_server.is_active():
            return

        self.object_center_x_callback_times += 1
        if self.object_center_x_callback_times == self.object_center_x_callback_times_threshold:
            self.object_center_x_callback_times = 0

            object_center_x = data.data
            rospy.logdebug('object center x: %s', object_center_x)

            # 计算物体中心与图像中点偏差
            bias_x = round((object_center_x - CAMERA_WIDTH/2)/(CAMERA_WIDTH/2), 3)

            # 判断是否在偏差阈值之内
            if bias_x < BIAS_THRESHOLD and bias_x > (-BIAS_THRESHOLD):
                bias_x = 0

                self.bias_x_correct_times += 1
                if self.bias_x_correct_times == self.bias_x_correct_times_threshold:
                    self.bias_x_correct_times = 0
                    self.flag = True

                    # 发布结果(result)
                    rospy.loginfo('Arm swing succeeded')
                    self._result.joint_states = self.joint_states
                    self._action_server.set_succeeded(self._result)

            # 比例控制信号计算
            control_x = CONTROL_RATIO * bias_x
            rospy.logdebug('control x: %s', control_x)

            # 调整关节状态
            self.joint_cmd_data_point.positions = self.joint_states
            self.joint_cmd_data_point.positions[0] = self.joint_states[0] - control_x
            rospy.logdebug('publish data is %s', self.joint_cmd_data_point.positions)
            self.joint_cmd_data.points = [self.joint_cmd_data_point]

            # 发布反馈(feedback)
            self._feedback.joint_states = self.joint_states
            self._action_server.publish_feedback(self._feedback)

            # 发送控制信号
            self.joint_cmd_pub.publish(self.joint_cmd_data)
            rospy.logdebug('publish data OK')

        rospy.logdebug('object center x callback OK')


if __name__ == '__main__':
    rospy.init_node('npu_robot_arm_swing_server', anonymous=True, log_level=envs.log_level)
    arm_swing_server = ArmSwingActionServer()
